{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/NeuromatchAcademy/course-content/blob/master/tutorials/W2D1_DeepLearning/W2D1_Tutorial2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Neuromatch Academy: Week 2, Day 1, Tutorial 2\n",
    "# Deep Learning: Convolutional Neural Networks\n",
    "\n",
    "**Content creators**: Jorge A. Menendez, Carsen Stringer \n",
    "\n",
    "**Content reviewers**: Roozbeh Farhoodi, Ella Batty, Kshitij Dwivedi, Spiros Chavlis, Michael Waskom\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "#Tutorial Objectives\n",
    "\n",
    "In this short tutorial, we'll go through an introduction to 2D convolutions and apply a convolutional network to an image to prepare for creating normative models in tutorial 3.\n",
    "\n",
    "In this tutorial, we will \n",
    "* Understand the basics of 2D convolution\n",
    "* Build a convolutional layer using PyTorch\n",
    "* Visualize and analyze its outputs\n",
    "* (Bonus) Build an encoding model from images to neural responses"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Setup\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:44.174186Z",
     "iopub.status.busy": "2021-05-31T20:48:44.173480Z",
     "iopub.status.idle": "2021-05-31T20:48:44.785651Z",
     "shell.execute_reply": "2021-05-31T20:48:44.785037Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "from matplotlib import pyplot as plt\n",
    "import matplotlib as mpl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:44.798386Z",
     "iopub.status.busy": "2021-05-31T20:48:44.797620Z",
     "iopub.status.idle": "2021-05-31T20:48:44.823192Z",
     "shell.execute_reply": "2021-05-31T20:48:44.823673Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Figure settings\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "plt.style.use(\"https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:44.829404Z",
     "iopub.status.busy": "2021-05-31T20:48:44.828843Z",
     "iopub.status.idle": "2021-05-31T20:48:45.840484Z",
     "shell.execute_reply": "2021-05-31T20:48:45.839900Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Data retrieval and loading\n",
    "import hashlib\n",
    "import requests\n",
    "\n",
    "fname = \"W3D4_stringer_oribinned6_split.npz\"\n",
    "url = \"https://osf.io/p3aeb/download\"\n",
    "expected_md5 = \"b3f7245c6221234a676b71a1f43c3bb5\"\n",
    "\n",
    "if not os.path.isfile(fname):\n",
    "  try:\n",
    "    r = requests.get(url)\n",
    "  except requests.ConnectionError:\n",
    "    print(\"!!! Failed to download data !!!\")\n",
    "  else:\n",
    "    if r.status_code != requests.codes.ok:\n",
    "      print(\"!!! Failed to download data !!!\")\n",
    "    elif hashlib.md5(r.content).hexdigest() != expected_md5:\n",
    "      print(\"!!! Data download appears corrupted !!!\")\n",
    "    else:\n",
    "      with open(fname, \"wb\") as fid:\n",
    "        fid.write(r.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:45.857021Z",
     "iopub.status.busy": "2021-05-31T20:48:45.854111Z",
     "iopub.status.idle": "2021-05-31T20:48:45.859261Z",
     "shell.execute_reply": "2021-05-31T20:48:45.858738Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Plotting Functions\n",
    "\n",
    "def plot_tuning(ax, stimuli, respi_train, respi_test, neuron_index, linewidth=2):\n",
    "  \"\"\"Plot the tuning curve of a neuron\"\"\"\n",
    "\n",
    "  ax.plot(stimuli, respi_train, 'y', linewidth=linewidth)  # plot its responses as a function of stimulus orientation\n",
    "  ax.plot(stimuli, respi_test, 'm', linewidth=linewidth)  # plot its responses as a function of stimulus orientation\n",
    "  ax.set_title('neuron %i' % neuron_index)\n",
    "  ax.set_xlabel('stimulus orientation ($^o$)')\n",
    "  ax.set_ylabel('neural response')\n",
    "  ax.set_xticks(np.linspace(0, 360, 5))\n",
    "  ax.set_ylim([-0.5, 2.4])\n",
    "\n",
    "\n",
    "def show_stimulus(img, ax=None):\n",
    "  \"\"\"Visualize a stimulus\"\"\"\n",
    "  if ax is None:\n",
    "    ax = plt.gca()\n",
    "  ax.imshow(img+0.5, cmap=mpl.cm.binary)\n",
    "  ax.set_xticks([])\n",
    "  ax.set_yticks([])\n",
    "  ax.spines['left'].set_visible(False)\n",
    "  ax.spines['bottom'].set_visible(False)\n",
    "\n",
    "\n",
    "def plot_example_activations(stimuli, act, channels=[0]):\n",
    "  \"\"\" plot activations act and corresponding stimulus\n",
    "  Args:\n",
    "        stimuli: stimulus input to convolutional layer (n x h x w) or (h x w)\n",
    "        act: activations of convolutional layer (n_bins x conv_channels x n_bins)\n",
    "        channels: which conv channels to plot\n",
    "  \"\"\"\n",
    "  if stimuli.ndim>2:\n",
    "    n_stimuli = stimuli.shape[0]\n",
    "  else:\n",
    "    stimuli = stimuli.unsqueeze(0)\n",
    "    n_stimuli = 1\n",
    "\n",
    "  fig, axs = plt.subplots(n_stimuli,1+len(channels),figsize=(12,12))\n",
    "\n",
    "  # plot stimulus\n",
    "  for i in range(n_stimuli):\n",
    "    show_stimulus(stimuli[i].squeeze(), ax=axs[i, 0])\n",
    "    axs[i, 0].set_title('stimulus')\n",
    "\n",
    "    # plot example activations\n",
    "    for k, (channel, ax) in enumerate(zip(channels, axs[i][1:])):\n",
    "      img=ax.imshow(act[i,channel], vmin=-6, vmax=6, cmap='bwr')\n",
    "      ax.set_xlabel('x-pos')\n",
    "      ax.set_ylabel('y-pos')\n",
    "      ax.set_title('channel %d'%channel)\n",
    "  ax = fig.add_axes([1.05,0.8,0.01,0.1])\n",
    "  plt.colorbar(img, cax=ax)\n",
    "  ax.set_title('activation\\n strength')\n",
    "\n",
    "\n",
    "def plot_weights(weights, channels=[0], colorbar=True):\n",
    "  \"\"\" plot convolutional channel weights\n",
    "  Args:\n",
    "        weights: weights of convolutional filters (conv_channels x K x K)\n",
    "        channels: which conv channels to plot\n",
    "  \"\"\"\n",
    "  wmax = torch.abs(weights).max()\n",
    "  fig, axs = plt.subplots(1,len(channels), figsize=(12,2.5))\n",
    "  for i, channel in enumerate(channels):\n",
    "    im = axs[i].imshow(weights[channel,0], vmin=-wmax, vmax=wmax, cmap='bwr')\n",
    "    axs[i].set_title('channel %d'%channel)\n",
    "\n",
    "  if colorbar:\n",
    "    ax = fig.add_axes([1, 0.1, 0.05, 0.8])\n",
    "    plt.colorbar(im, ax=ax)\n",
    "    ax.axis('off')\n",
    "\n",
    "\n",
    "def plot_prediction(ax, y_pred, y_train, y_test):\n",
    "  \"\"\" plot prediction of neural response + test neural response \"\"\"\n",
    "  ax.plot(y_train, 'y', linewidth=1)\n",
    "  ax.plot(y_test,color='m')\n",
    "  ax.plot(y_pred, 'g', linewidth=3)\n",
    "  ax.set_xlabel('stimulus bin')\n",
    "  ax.set_ylabel('response')\n",
    "\n",
    "\n",
    "def plot_training_curves(train_loss, test_loss):\n",
    "  \"\"\"\n",
    "  Args:\n",
    "    train_loss (list): training error over iterations\n",
    "    test_loss (list): n_test x 1 tensor with orientations of the\n",
    "      stimuli corresponding to each row of train_data, in radians\n",
    "    predicted_test_labels (torch.Tensor): n_test x 1 tensor with predicted orientations of the\n",
    "      stimuli from decoding neural network\n",
    "\n",
    "  \"\"\"\n",
    "\n",
    "  f, ax = plt.subplots()\n",
    "  # Plot the training loss over iterations of GD\n",
    "  ax.plot(train_loss)\n",
    "  # Plot the testing loss over iterations of GD\n",
    "  ax.plot(test_loss, '.', markersize=10)\n",
    "  ax.legend(['train loss', 'test loss'])\n",
    "  ax.set(xlabel=\"Gradient descent iteration\", ylabel=\"Mean squared error\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:45.877590Z",
     "iopub.status.busy": "2021-05-31T20:48:45.872273Z",
     "iopub.status.idle": "2021-05-31T20:48:45.885942Z",
     "shell.execute_reply": "2021-05-31T20:48:45.885397Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Helper Functions\n",
    "\n",
    "def load_data_split(data_name=fname):\n",
    "  \"\"\"Load mouse V1 data from Stringer et al. (2019)\n",
    "\n",
    "  Data from study reported in this preprint:\n",
    "  https://www.biorxiv.org/content/10.1101/679324v2.abstract\n",
    "\n",
    "  These data comprise time-averaged responses of ~20,000 neurons\n",
    "  to ~4,000 stimulus gratings of different orientations, recorded\n",
    "  through Calcium imaginge. The responses have been normalized by\n",
    "  spontaneous levels of activity and then z-scored over stimuli, so\n",
    "  expect negative numbers. The repsonses were split into train and\n",
    "  test and then each set were averaged in bins of 6 degrees.\n",
    "\n",
    "  This function returns the relevant data (neural responses and\n",
    "  stimulus orientations) in a torch.Tensor of data type torch.float32\n",
    "  in order to match the default data type for nn.Parameters in\n",
    "  Google Colab.\n",
    "\n",
    "  It will hold out some of the trials when averaging to allow us to have test\n",
    "  tuning curves.\n",
    "\n",
    "  Args:\n",
    "    data_name (str): filename to load\n",
    "\n",
    "  Returns:\n",
    "    resp_train (torch.Tensor): n_stimuli x n_neurons matrix of neural responses,\n",
    "        each row contains the responses of each neuron to a given stimulus.\n",
    "        As mentioned above, neural \"response\" is actually an average over\n",
    "        responses to stimuli with similar angles falling within specified bins.\n",
    "    resp_test (torch.Tensor): n_stimuli x n_neurons matrix of neural responses,\n",
    "        each row contains the responses of each neuron to a given stimulus.\n",
    "        As mentioned above, neural \"response\" is actually an average over\n",
    "        responses to stimuli with similar angles falling within specified bins\n",
    "    stimuli: (torch.Tensor): n_stimuli x 1 column vector with orientation\n",
    "        of each stimulus, in degrees. This is actually the mean orientation\n",
    "        of all stimuli in each bin.\n",
    "\n",
    "  \"\"\"\n",
    "  with np.load(data_name) as dobj:\n",
    "    data = dict(**dobj)\n",
    "  resp_train = data['resp_train']\n",
    "  resp_test = data['resp_test']\n",
    "  stimuli = data['stimuli']\n",
    "\n",
    "  # Return as torch.Tensor\n",
    "  resp_train_tensor = torch.tensor(resp_train, dtype=torch.float32)\n",
    "  resp_test_tensor = torch.tensor(resp_test, dtype=torch.float32)\n",
    "  stimuli_tensor = torch.tensor(stimuli, dtype=torch.float32)\n",
    "\n",
    "  return resp_train_tensor, resp_test_tensor, stimuli_tensor\n",
    "\n",
    "\n",
    "def grating(angle, sf=1 / 28, res=0.1, patch=False):\n",
    "  \"\"\"Generate oriented grating stimulus\n",
    "\n",
    "  Args:\n",
    "    angle (float): orientation of grating (angle from vertical), in degrees\n",
    "    sf (float): controls spatial frequency of the grating\n",
    "    res (float): resolution of image. Smaller values will make the image\n",
    "      smaller in terms of pixels. res=1.0 corresponds to 640 x 480 pixels.\n",
    "    patch (boolean): set to True to make the grating a localized\n",
    "      patch on the left side of the image. If False, then the\n",
    "      grating occupies the full image.\n",
    "\n",
    "  Returns:\n",
    "    torch.Tensor: (res * 480) x (res * 640) pixel oriented grating image\n",
    "\n",
    "  \"\"\"\n",
    "\n",
    "  angle = np.deg2rad(angle)  # transform to radians\n",
    "\n",
    "  wpix, hpix = 640, 480  # width and height of image in pixels for res=1.0\n",
    "\n",
    "  xx, yy = np.meshgrid(sf * np.arange(0, wpix * res) / res, sf * np.arange(0, hpix * res) / res)\n",
    "\n",
    "  if patch:\n",
    "    gratings = np.cos(xx * np.cos(angle + .1) + yy * np.sin(angle + .1))  # phase shift to make it better fit within patch\n",
    "    gratings[gratings < 0] = 0\n",
    "    gratings[gratings > 0] = 1\n",
    "    xcent = gratings.shape[1] * .75\n",
    "    ycent = gratings.shape[0] / 2\n",
    "    xxc, yyc = np.meshgrid(np.arange(0, gratings.shape[1]), np.arange(0, gratings.shape[0]))\n",
    "    icirc = ((xxc - xcent) ** 2 + (yyc - ycent) ** 2) ** 0.5 < wpix / 3 / 2 * res\n",
    "    gratings[~icirc] = 0.5\n",
    "\n",
    "  else:\n",
    "    gratings = np.cos(xx * np.cos(angle) + yy * np.sin(angle))\n",
    "    gratings[gratings < 0] = 0\n",
    "    gratings[gratings > 0] = 1\n",
    "\n",
    "  gratings -= 0.5\n",
    "\n",
    "  # Return torch tensor\n",
    "  return torch.tensor(gratings, dtype=torch.float32)\n",
    "\n",
    "\n",
    "def filters(K=7):\n",
    "  \"\"\" make example filters, some center-surround and gabors\n",
    "  Returns:\n",
    "      filters: 6 x K x K\n",
    "  \"\"\"\n",
    "  grid = np.linspace(-K/2, K/2, K).astype(np.float32)\n",
    "  xx,yy = np.meshgrid(grid, grid, indexing='ij')\n",
    "\n",
    "  # create center-surround filters\n",
    "  sigma = 1.1\n",
    "  gaussian = np.exp(-(xx**2 + yy**2)**0.5/(2*sigma**2))\n",
    "  wide_gaussian = np.exp(-(xx**2 + yy**2)**0.5/(2*(sigma*2)**2))\n",
    "  center_surround = gaussian - 0.5 * wide_gaussian\n",
    "\n",
    "  # create gabor filters\n",
    "  thetas = np.array([0, 45, 90, 135]) * np.pi/180\n",
    "  gabors = np.zeros((len(thetas), K, K), np.float32)\n",
    "  lam = 10\n",
    "  phi = np.pi/2\n",
    "  gaussian = np.exp(-(xx**2 + yy**2)**0.5/(2*(sigma*0.4)**2))\n",
    "  for i,theta in enumerate(thetas):\n",
    "    x = xx*np.cos(theta) + yy*np.sin(theta)\n",
    "    gabors[i] = gaussian * np.cos(2*np.pi*x/lam + phi)\n",
    "\n",
    "  filters = np.concatenate((center_surround[np.newaxis,:,:],\n",
    "                            -1*center_surround[np.newaxis,:,:],\n",
    "                            gabors),\n",
    "                           axis=0)\n",
    "  filters /= np.abs(filters).max(axis=(1,2))[:,np.newaxis,np.newaxis]\n",
    "  # convert to torch\n",
    "  filters = torch.from_numpy(filters)\n",
    "  # add channel axis\n",
    "  filters = filters.unsqueeze(1)\n",
    "\n",
    "  return filters\n",
    "\n",
    "\n",
    "def regularized_MSE_loss(output, target, weights=None, L2_penalty=0, L1_penalty=0):\n",
    "  \"\"\"loss function for MSE\n",
    "\n",
    "  Args:\n",
    "    output (torch.Tensor): output of network\n",
    "    target (torch.Tensor): neural response network is trying to predict\n",
    "    weights (torch.Tensor): fully-connected layer weights of network (net.out_layer.weight)\n",
    "    L2_penalty : scaling factor of sum of squared weights\n",
    "    L1_penalty : scalaing factor for sum of absolute weights\n",
    "\n",
    "  Returns:\n",
    "    (torch.Tensor) mean-squared error with L1 and L2 penalties added\n",
    "\n",
    "  \"\"\"\n",
    "\n",
    "  loss_fn = nn.MSELoss()\n",
    "  loss = loss_fn(output, target)\n",
    "\n",
    "  if weights is not None:\n",
    "    L2 = L2_penalty * torch.square(weights).sum()\n",
    "    L1 = L1_penalty * torch.abs(weights).sum()\n",
    "    loss += L1 + L2\n",
    "\n",
    "  return loss\n",
    "\n",
    "def train(net, custom_loss, train_data, train_labels,\n",
    "          test_data=None, test_labels=None,\n",
    "          learning_rate=10, n_iter=500, L2_penalty=0., L1_penalty=0.):\n",
    "  \"\"\"Run gradient descent for network without batches\n",
    "\n",
    "  Args:\n",
    "    net (nn.Module): deep network whose parameters to optimize with SGD\n",
    "    custom_loss: loss function for network\n",
    "    train_data: training data (n_train x input features)\n",
    "    train_labels: training labels (n_train x output features)\n",
    "    test_data: test data (n_train x input features)\n",
    "    test_labels: test labels (n_train x output features)\n",
    "    learning_rate (float): learning rate for gradient descent\n",
    "    n_epochs (int): number of epochs to run gradient descent\n",
    "    L2_penalty (float): magnitude of L2 penalty\n",
    "    L1_penalty (float): magnitude of L1 penalty\n",
    "\n",
    "  Returns:\n",
    "    train_loss: training loss across iterations\n",
    "    test_loss: testing loss across iterations\n",
    "\n",
    "  \"\"\"\n",
    "  optimizer = optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9) # Initialize PyTorch SGD optimizer\n",
    "  train_loss = np.nan * np.zeros(n_iter)  # Placeholder for train loss\n",
    "  test_loss = np.nan * np.zeros(n_iter)  # Placeholder for test loss\n",
    "\n",
    "  # Loop over epochs\n",
    "  for i in range(n_iter):\n",
    "    if n_iter < 10:\n",
    "      for param_group in self.optimizer.param_groups:\n",
    "        param_group['lr'] = np.linspace(0, learning_rate, 10)[n_iter]\n",
    "    y_pred = net(train_data) # Forward pass: compute predicted y by passing train_data to the model.\n",
    "\n",
    "    if L2_penalty>0 or L1_penalty>0:\n",
    "      weights = net.out_layer.weight\n",
    "      loss = custom_loss(y_pred, train_labels, weights, L2_penalty, L1_penalty)\n",
    "    else:\n",
    "      loss = custom_loss(y_pred, train_labels)\n",
    "\n",
    "    ### Update parameters\n",
    "    optimizer.zero_grad() # zero out gradients\n",
    "    loss.backward() # Backward pass: compute gradient of the loss with respect to model parameters\n",
    "    optimizer.step() # step parameters in gradient direction\n",
    "    train_loss[i] = loss.item()  # .item() transforms the tensor to a scalar and does .detach() for us\n",
    "\n",
    "    # Track progress\n",
    "    if (i+1) % (n_iter // 10) == 0 or i==0:\n",
    "      if test_data is not None and test_labels is not None:\n",
    "        y_pred = net(test_data)\n",
    "        if L2_penalty>0 or L1_penalty>0:\n",
    "          loss = custom_loss(y_pred, test_labels, weights, L2_penalty, L1_penalty)\n",
    "        else:\n",
    "          loss = custom_loss(y_pred, test_labels)\n",
    "        test_loss[i] = loss.item()\n",
    "        print(f'iteration {i+1}/{n_iter} | train loss: {train_loss[i]:.4f} | test loss: {test_loss[i]:.4f}')\n",
    "      else:\n",
    "        print(f'iteration {i+1}/{n_iter} | train loss: {train_loss[i]:.4f}')\n",
    "\n",
    "  return train_loss, test_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Section 1: Introduction to 2D convolutions\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 1.1: What is a 2D convolution?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A 2D convolution is an integral of the product of a filter $f$ and an input image $I$ computed at various positions as the filter is slid across the input. The output of the convolution operation at position $(x,y)$ can be written as follows, where the filter $f$ is size $(K,K)$:\n",
    "\\begin{equation}\n",
    "C(x,y) = \\sum_{k_x=-K/2}^{K/2} \\sum_{k_y=-K/2}^{K/2} f(k_x,k_y) I(x+k_x,y+k_y)\n",
    "\\end{equation}\n",
    "This **convolutional filter** is often called a **kernel**.\n",
    "\n",
    "Here is an illustration of a 2D convolution from this [article](https://towardsdatascience.com/light-on-math-machine-learning-intuitive-guide-to-convolution-neural-networks-e3f054dd5daa):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 326
    },
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:45.892650Z",
     "iopub.status.busy": "2021-05-31T20:48:45.892112Z",
     "iopub.status.idle": "2021-05-31T20:48:45.894994Z",
     "shell.execute_reply": "2021-05-31T20:48:45.895462Z"
    },
    "outputId": "94cc3918-b79d-47bb-c8cb-3ea7480667f1"
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "#@markdown Execute this cell to view convolution gif\n",
    "\n",
    "from IPython.display import Image\n",
    "Image(url='https://miro.medium.com/max/700/1*5BwZUqAqFFP5f3wKYQ6wJg.gif')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 1.2: 2D convolutions in deep learning\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 517
    },
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:45.902036Z",
     "iopub.status.busy": "2021-05-31T20:48:45.901372Z",
     "iopub.status.idle": "2021-05-31T20:48:45.928033Z",
     "shell.execute_reply": "2021-05-31T20:48:45.928606Z"
    },
    "outputId": "521ef866-fa7b-4022-a688-0794a32ecbf0"
   },
   "outputs": [],
   "source": [
    "#@title Video 1: 2D Convolutions\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"igLkHuuJurE\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "In the [intro](https://www.youtube.com/watch?v=IZvcy0Myb3M) you learned about convolutional networks for computer vision tasks and applications to neuroscience. Object recognition was essentially an unsolved problem in machine learning until the [advent](https://en.wikipedia.org/wiki/AlexNet) of techniques for effectively training *deep* convolutional neural networks. See Bonus Section 1 for more information on why we use CNNs to model the visual system.\n",
    "\n",
    "Convolutional neural networks consist of 2D convolutional layers, RELU non-linearities, 2D pooling layers, and at the output, a fully connected layer. We will see an example network with all these components in tutorial 3. \n",
    "\n",
    "A 2D convolutional layer consists of multiple channels. Each **channel** is a 2D convolutional filter applied to the input resulting in several units, where the number of units depends on the *stride* you set. In most applications, especially with small filter sizes, a stride of 1 is used.\n",
    "\n",
    "(*Technical note*: if filter size *K* is odd and you set the *pad=K//2* and *stride=1* (as is shown below), you get a **channel** of units that is the same size as the input.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 470
    },
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:45.934296Z",
     "iopub.status.busy": "2021-05-31T20:48:45.933288Z",
     "iopub.status.idle": "2021-05-31T20:48:45.936041Z",
     "shell.execute_reply": "2021-05-31T20:48:45.936480Z"
    },
    "outputId": "c3b74a90-326c-48d8-a4e5-b0d4372f0882"
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "#@markdown Execute this cell to view convolution gif\n",
    "\n",
    "from IPython.display import Image\n",
    "Image(url='https://miro.medium.com/max/790/1*1okwhewf5KCtIPaFib4XaA.gif')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 1.3: 2D convolutions in Pytorch\n",
    "\n",
    "In Tutorial 1, fully connected linear layers were used to decode stimuli from neural activity. Convolutional layers are different from their fully connected counterparts in two ways:\n",
    "  * In a fully connected layer, each unit computes a weighted sum over all the input units. In a convolutional layer, on the other hand, each unit computes a weighted sum over only a small patch of the input, referred to as the unit's **receptive field**. When the input is an image, the receptive field can be thought of as a local patch of pixels.\n",
    "  * In a fully connected layer, each unit uses its own independent set of weights to compute the weighted sum. In a convolutional layer, all the units (within the same channel) **share the same weights**. This set of shared weights is called the **convolutional filter or kernel**. The result of this computation is a convolution, where each unit has computed the same weighted sum over a different part of the input. This reduces the number of parameters in the network substantially.\n",
    "\n",
    "<p align=\"center\">\n",
    "  <img src=\"https://github.com/NeuromatchAcademy/course-content/blob/master/tutorials/static/weight-sharing.png?raw=true\" width=\"700\" />\n",
    "</p>\n",
    "\n",
    "\n",
    "We will compute the difference in the number of weights for a fully connected layer versus a convolutional layer in the Think exercise below.\n",
    "\n",
    "First, let's visualize the stimuli in the dataset from tutorial 1. During the neural recordings from [Stringer et al 2019](https://www.biorxiv.org/content/10.1101/679324v2.abstract), mice were presented oriented gratings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 215
    },
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:45.974598Z",
     "iopub.status.busy": "2021-05-31T20:48:45.959998Z",
     "iopub.status.idle": "2021-05-31T20:48:46.432810Z",
     "shell.execute_reply": "2021-05-31T20:48:46.432343Z"
    },
    "outputId": "c2817d07-daed-4a2b-bc87-4991f160f296"
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "#@markdown Execute this cell to plot example stimuli\n",
    "\n",
    "orientations = np.linspace(-90, 90, 5)\n",
    "\n",
    "h = 3\n",
    "n_col = len(orientations)\n",
    "fig, axs = plt.subplots(1, n_col, figsize=(h * n_col, h))\n",
    "\n",
    "h, w  = grating(0).shape  # height and width of stimulus\n",
    "print('stimulus size: %i x %i' % (h, w))\n",
    "\n",
    "for i, ori in enumerate(orientations):\n",
    "  stimulus = grating(ori)\n",
    "  axs[i].set_title(f'{ori: .0f}$^o$')\n",
    "  show_stimulus(stimulus, axs[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's implement 2D convolutional operations. We will use mutliple convolutional channels and implement this operation efficiently using pytorch. A *layer* of convolutional channels can be implemented with one line of code using the PyTorch class `nn.Conv2d()`, which requires the following arguments for initialization (see full documentation [here](https://pytorch.org/docs/master/generated/torch.nn.Conv2d.html)):\n",
    "  * $C^{in}$: the number of input channels\n",
    "  * $C^{out}$: the number of output channels (number of different convolutional filters)\n",
    "  * $K$: the size of the $C^{out}$ different convolutional filters\n",
    "  \n",
    "When you run the network, you can input a stimulus of arbitrary size $(H^{in}, W^{in})$, but it needs to be shaped as a 4D input $(N, C^{in}, H^{in}, W^{in})$, where $N$ is the number of images. In our case, $C^{in}=1$ because there is only one color channel (our images are grayscale, but often $C^{in}=3$ in image processing)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:46.439160Z",
     "iopub.status.busy": "2021-05-31T20:48:46.438522Z",
     "iopub.status.idle": "2021-05-31T20:48:46.442185Z",
     "shell.execute_reply": "2021-05-31T20:48:46.441687Z"
    }
   },
   "outputs": [],
   "source": [
    "class ConvolutionalLayer(nn.Module):\n",
    "  \"\"\"Deep network with one convolutional layer\n",
    "     Attributes: conv (nn.Conv2d): convolutional layer\n",
    "  \"\"\"\n",
    "  def __init__(self, c_in=1, c_out=6, K=7, filters=None):\n",
    "    \"\"\"Initialize layer\n",
    "\n",
    "    Args:\n",
    "        c_in: number of input stimulus channels\n",
    "        c_out: number of output convolutional channels\n",
    "        K: size of each convolutional filter\n",
    "        filters: (optional) initialize the convolutional weights\n",
    "\n",
    "    \"\"\"\n",
    "    super().__init__()\n",
    "    self.conv = nn.Conv2d(c_in, c_out, kernel_size=K,\n",
    "                          padding=K//2, stride=1)\n",
    "    if filters is not None:\n",
    "      self.conv.weight = nn.Parameter(filters)\n",
    "      self.conv.bias = nn.Parameter(torch.zeros((c_out,), dtype=torch.float32))\n",
    "\n",
    "  def forward(self, s):\n",
    "    \"\"\"Run stimulus through convolutional layer\n",
    "\n",
    "    Args:\n",
    "        s (torch.Tensor): n_stimuli x h x w tensor with stimuli\n",
    "\n",
    "    Returns:\n",
    "        (torch.Tensor): n_stimuli x c_out x w x h tensor with convolutional layer unit activations.\n",
    "\n",
    "    \"\"\"\n",
    "    s = s.unsqueeze(1) # add channel dimension\n",
    "    a = self.conv(s)  # output of convolutional layer\n",
    "\n",
    "    return a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Coding Exercise 1.2: 2D convolution in pytorch \n",
    "\n",
    "We will now run the convolutional layer on our stimulus. We will use gratings stimuli made using the function `grating`, which returns a stimulus which is 48 x 64. \n",
    "\n",
    "Reminder, `nn.Conv2d` takes in a tensor of size $(N, C^{in}, H^{in}, W^{in}$) where $N$ is the number of stimuli, $C^{in}$ is the number of input channels, and $(H^{in}, W^{in})$ is the size of the stimulus. We will need to add these first two dimensions to our stimulus, then input it to the convolutional layer.\n",
    "\n",
    "We will plot the outputs of the convolution. `convout` is a tensor of size $(N, C^{out}, H^{in}, W^{in})$ where $N$ is the number of examples and $C^{out}$ are the number of convolutional channels. It is the same size as the input because we used a stride of 1 and padding that is half the kernel size.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:46.447149Z",
     "iopub.status.busy": "2021-05-31T20:48:46.446644Z",
     "iopub.status.idle": "2021-05-31T20:48:46.450298Z",
     "shell.execute_reply": "2021-05-31T20:48:46.449801Z"
    }
   },
   "outputs": [],
   "source": [
    "# Convolution layer parameters\n",
    "K = 7 # filter size\n",
    "conv_channels = 6 # how many convolutional channels to have in our layer\n",
    "convout = np.zeros(0) # assign convolutional activations to convout\n",
    "\n",
    "################################################################################\n",
    "## TODO for students: create convolutional layer in pytorch\n",
    "# Complete and uncomment\n",
    "################################################################################\n",
    "\n",
    "### Initialize conv layer and add weights from function filters\n",
    "# you need to specify :\n",
    "# * the number of input channels c_in\n",
    "# * the number of output channels c_out\n",
    "# * the filter size K\n",
    "# convLayer = ConvolutionalLayer(..., filters=filters(K))\n",
    "\n",
    "### Create stimuli (H_in, W_in)\n",
    "stimuli = torch.zeros((len(orientations), 48, 64), dtype=torch.float32)\n",
    "orientations = [-90, -45, 0, 45, 90]\n",
    "for i,ori in enumerate(orientations):\n",
    "  stimuli[i] = grating(ori)\n",
    "\n",
    "# convout = convLayer(...)\n",
    "# convout = convout.detach() # detach gradients\n",
    "\n",
    "# plot_example_activations(stimuli, convout, channels=np.arange(0, conv_channels))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 912
    },
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:46.456124Z",
     "iopub.status.busy": "2021-05-31T20:48:46.455548Z",
     "iopub.status.idle": "2021-05-31T20:48:49.347903Z",
     "shell.execute_reply": "2021-05-31T20:48:49.348345Z"
    },
    "outputId": "6ac1dd1c-fccb-42cd-aee9-c900fecf786c"
   },
   "outputs": [],
   "source": [
    "# to_remove solution\n",
    "\n",
    "# Convolution layer parameters\n",
    "K = 7 # filter size\n",
    "conv_channels = 6 # how many convolutional channels to have in our layer\n",
    "\n",
    "convout = np.zeros(0) # assign convolutional activations to convout\n",
    "\n",
    "### Initialize conv layer with weights from function filters\n",
    "# you need to specify :\n",
    "# * the number of input channels c_in\n",
    "# * the number of output channels c_out\n",
    "# * the filter size K\n",
    "convLayer = ConvolutionalLayer(c_in=1, c_out=conv_channels, K=K,\n",
    "                               filters=filters(K))\n",
    "\n",
    "### Create stimuli (H_in, W_in)\n",
    "stimuli = torch.zeros((len(orientations), 48, 64), dtype=torch.float32)\n",
    "orientations = [-90, -45, 0, 45, 90]\n",
    "for i,ori in enumerate(orientations):\n",
    "  stimuli[i] = grating(ori)\n",
    "\n",
    "convout = convLayer(stimuli)\n",
    "convout = convout.detach() # detach gradients\n",
    "\n",
    "with plt.xkcd():\n",
    "  plot_example_activations(stimuli, convout, channels=np.arange(0, conv_channels))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Think! 2.3: Output and weight shapes in a convolutional layer\n",
    "\n",
    "Let's think about the shape of the weights and outputs of `convLayer`:\n",
    "  - How many convolutional activations are there in each channel, and why is it this size?\n",
    "  - How many weights does `convLayer` have? \n",
    "  - How many weights would it have if it were a fully connected layer?\n",
    "\n",
    "Additionally, let's think about why the activations look the way they do. It seems like for all channels the activations are only non-zero for edges of the gratings (where the grating goes from white-to-black and from black-to-white).\n",
    "  -  Channel 0 and 1 seem to respond to every edge regardless of the orientation, but their signs are different. What type of filter might produce these types of responses?\n",
    "  - Channels 2-5 seem to respond differently depending on the orientation of the stimulus. What type of filter might produce these types of responses?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:49.353361Z",
     "iopub.status.busy": "2021-05-31T20:48:49.352782Z",
     "iopub.status.idle": "2021-05-31T20:48:49.354976Z",
     "shell.execute_reply": "2021-05-31T20:48:49.355442Z"
    }
   },
   "outputs": [],
   "source": [
    "# to_remove explanation\n",
    "\n",
    "\"\"\"\n",
    "1. There are H * W activations in each channel because the stride in the\n",
    "convolutional layer is set to 1, and the padding is set to K//2.\n",
    "2. The convLayer has K * K * C_out weights and C_out bias terms.\n",
    "3. A fully connected layer would have (H * W) * C_out weights and C_out bias terms.\n",
    "\n",
    "4. A center-surround filter will respond to a change in luminance (black-to-white\n",
    "or white-to-black) regardless of orientation.\n",
    "5. A gabor filter of different orientations will respond this way. See the\n",
    "exercise below for more explanation.\n",
    "\"\"\";"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Please see Bonus Section 2 to visualize the convolutional filter weights. See Bonus Section 3 to use CNNs as encoding models of neurons (by fitting directly to neural responses)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Summary\n",
    "\n",
    "In this notebook, we built a 2D convolutional layer which is meant to represent the responses of neurons in mouse visual cortex, or the responses of neurons \n",
    "which are inputs to mouse visual cortex.\n",
    "\n",
    "In Tutorial 3, we will add to this 2D convolutional layer a fully-connected layer and train this model to predict whether an orientation is left or right. We will see if the convolutional filters it learns are similar to mouse visual cortex."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Bonus"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus Section 1: Why CNN's?\n",
    "\n",
    "CNN models are particularly [well-suited](https://www.nature.com/articles/nn.4244) to modeling the visual system for a number of reasons:\n",
    "\n",
    "1. **Distributed computation**: like any other neural network, CNN's use distributed representations to compute -- much like the brain seems to do. Such models, therefore, provide us with a vocabulary with which to talk and think about such distributed representations. Because we know the exact function the model is built to perform (e.g. orientation discrimination), we can analyze its internal representations with respect to this function and begin to interpret why the representations look the way they do. Most importantly, we can then use these insights to analyze the structure of neural representations observed in recorded population activity. We can qualitatively and quantitatively compare the representations we see in the model and in a given population of real neurons to hopefully tease out the computations it performs.\n",
    "\n",
    "2. **Hierarchical architecture**: like in any other deep learning architecture, each layer of a deep CNN comprises a non-linear transformation of the previous layer. Thus, there is a natural hierarchy whereby layers closer to the network output represent increasingly more abstract information about the input image. For example, in a network trained to do object recognition, the early layers might represent information about edges in the image, whereas later layers closer to the output might represent various object categories. This resembles the [hierarchical structure of the visual system](https://pubmed.ncbi.nlm.nih.gov/1822724/), where [lower-level areas](https://www.jneurosci.org/content/25/46/10577.short) (e.g. retina, V1) represent visual features of the sensory input and [higher-level areas](https://www.sciencedirect.com/science/article/pii/S089662731200092X) (e.g. V4, IT) represent properties of objects in the visual scene. We can then naturally use a single CNN to model multiple visual areas, using early CNN layers to model lower-level visual areas and late CNN layers to model higher-level visual areas.\n",
    "  \n",
    "  Relative to fully connected networks, CNN's, in fact, have further hierarchical structure built-in through the max pooling layers. Recall that each output of a convolution + pooling block is the result of processing a local patch of the inputs to that block. If we stack such blocks in a sequence, then the outputs of each block will be sensitive to increasingly larger regions of the initial raw input to the network: an output from the first block is sensitive to a single patch of these inputs, corresponding to its \"receptive field\"; an output from the second block is sensitive to a patch of outputs from the first block, which together are sensitive to a larger patch of raw inputs comprising the union of their receptive fields. Receptive fields thus get larger for deeper layers (see [here](http://colah.github.io/posts/2014-07-Conv-Nets-Modular/) for a nice visual depiction of this). This resembles primate visual systems, where neurons in higher-level visual areas respond to stimuli in wider regions of the visual field than neurons in lower-level visual areas.\n",
    "\n",
    "3. **Convolutional layers**: through the weight sharing constraint, the outputs of each channel of a convolutional layer process different parts of the input image in exactly the same way. This architectural constraint effectively builds into the network the assumption that objects in the world typically look the same regardless of where they are in space. This is useful for modeling the visual system for two (largely separate) reasons:\n",
    "  * Firstly, this assumption is generally valid in mammalian visual systems, since mammals tend to view the same object from many perspectives. Two neurons at a similar hierarchy in the visual system with different receptive fields could thus end up receiving statistically similar synaptic inputs, so that the synaptic weights developed over time may end up being similar as well.\n",
    "  * Secondly, this architecture significantly improves object recognition ability. Object recognition was essentially an unsolved problem in machine learning until the [advent](https://en.wikipedia.org/wiki/AlexNet) of techniques for effectively training *deep* convolutional neural networks. Fully connected networks on their own can't achieve object recognition abilities anywhere close to human levels, making them bad models of human object recognition. Indeed, it is generally the case that [the better a neural network model is at object recognition, the closer the match between its representations and those observed in the brain](https://www.pnas.org/content/111/23/8619.short). That said, it is worth noting that our much simpler orientation discrimination task (in Tutorial 3) can be solved by relatively simple networks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Bonus Section 2: Understanding activations from weight"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bonus Coding Exercise 2: Visualizing convolutional filter weights\n",
    "\n",
    "Why do the activations look the way they do? Let's look at the weights (`convLayer.conv.weight`) of the convolutional filters and try to interpret them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:49.359753Z",
     "iopub.status.busy": "2021-05-31T20:48:49.358626Z",
     "iopub.status.idle": "2021-05-31T20:48:49.360493Z",
     "shell.execute_reply": "2021-05-31T20:48:49.360950Z"
    }
   },
   "outputs": [],
   "source": [
    "################################################################################\n",
    "## TODO for students: create convolutional layer in pytorch\n",
    "# Complete and uncomment\n",
    "################################################################################\n",
    "\n",
    "# get weights of conv layer in convLayer\n",
    "# weights = ...\n",
    "# print(weights.shape) # can you identify what each of the dimensions are?\n",
    "\n",
    "# plot_weights(weights, channels=np.arange(0, conv_channels))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 318
    },
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:49.365536Z",
     "iopub.status.busy": "2021-05-31T20:48:49.364988Z",
     "iopub.status.idle": "2021-05-31T20:48:49.997873Z",
     "shell.execute_reply": "2021-05-31T20:48:49.997367Z"
    },
    "outputId": "1e804e43-d1b1-4f2a-fdda-7fef50e37d7c"
   },
   "outputs": [],
   "source": [
    "# to_remove solution\n",
    "\n",
    "# get weights of conv layer in convLayer\n",
    "weights = convLayer.conv.weight.detach()\n",
    "print(weights.shape) # can you identify what each of the dimensions are?\n",
    "\n",
    "with plt.xkcd():\n",
    "  plot_weights(weights, channels=np.arange(0, conv_channels))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the function `filters` we premade center-surround filters and Gabor filters of various orientations. [Gabor filters](https://en.wikipedia.org/wiki/Gabor_filter) have a positive region next to a negative region and the orientation of these regions of the filter determine the orientation of edges to which they respond. \n",
    "\n",
    "In visual cortex, Hubel and Wiesel discovered simple cells, which would correspond to a unit in the channels 2-5 above with Gabor filters. There are also some neurons with activity that resembles center-surround filters, which would correspond to the first two convolutional channels above.\n",
    "\n",
    "There were additional cells discovered by Hubel and Wiesel - complex cells - that respond to an oriented grating regardless of where the bars are exactly (note that the responses we see are specific to where the bars are). These cells therefore have some level of translation invariance. This is something that convolutional neural networks try to replicate -- e.g. a grating is still oriented horizontal even if it moves slightly, and a cat is still a cat even if it's in a different position in the image."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Think!\n",
    "  - How might you create a complex cell and have responses that are translation invariant?\n",
    "  - How might you create a cell that responds to multiple orientations?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:50.002413Z",
     "iopub.status.busy": "2021-05-31T20:48:50.001886Z",
     "iopub.status.idle": "2021-05-31T20:48:50.007264Z",
     "shell.execute_reply": "2021-05-31T20:48:50.006728Z"
    }
   },
   "outputs": [],
   "source": [
    "# to_remove explanation\n",
    "\n",
    "\"\"\"\n",
    "1. A complex cell combines two rectified Gabor filters of the same orientation but with\n",
    "different phases. This is something a neural network can reproduce with a\n",
    "RELU layer and combinations of convolutional channels. Additionally neural networks\n",
    "employ a strategy called pooling layers, which combine units in a small region and\n",
    "take the maximum value of their activations as the output. We will see these\n",
    "layers in the next tutorial, where we will try to decode orientation directly from\n",
    "images.\n",
    "2. A cell that responds to multiple orientations would have to be a sum of multiple\n",
    "convolutional channels. In the next bonus section, we will hook up this convolutional\n",
    "layer to a fully connected layer.\n",
    "\"\"\";"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Bonus Section 3: Neural encoding model\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 517
    },
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:50.012812Z",
     "iopub.status.busy": "2021-05-31T20:48:50.011909Z",
     "iopub.status.idle": "2021-05-31T20:48:50.035527Z",
     "shell.execute_reply": "2021-05-31T20:48:50.034986Z"
    },
    "outputId": "e482918b-9e71-4fa5-cdf9-8ad718c174b9"
   },
   "outputs": [],
   "source": [
    "#@title Video 2: Convolutional Encoding Model\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"TCAQYIoqcgA\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "In neuroscience, we often want to understand how the brain represents external stimuli. One approach to discovering these representations is to build an encoding model that takes as input the external stimuli (in this case grating stimuli) and outputs the neural responses. \n",
    "\n",
    "Because visual cortex is often thought to be a convolutional network where the same filters are combined across the visual field, we will use a model with a convolutional layer. We learned how to build a convolutional layer in the previous section. We will add to this convolutional layer a fully connected layer from the output of the convolutions to the neurons. We will then visualize the weights of this fully connected layer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bonus Section 3.1: Neural tuning curves\n",
    "\n",
    "In the next cell, we plot the turning curves of a random subset of neurons. We have binned the stimuli orientations more than in Tutorial 1. We create the gratings images as above for the 60 orientations below, and save them to a variable `grating_stimuli`.\n",
    "\n",
    "Rerun the cell to look at different example neurons and observe the diversity of tuning curves in the population. How can we fit these neural responses with an encoding model?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 537
    },
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:50.043004Z",
     "iopub.status.busy": "2021-05-31T20:48:50.042472Z",
     "iopub.status.idle": "2021-05-31T20:48:51.738374Z",
     "shell.execute_reply": "2021-05-31T20:48:51.738853Z"
    },
    "outputId": "791505c8-5d1c-4f04-f080-95cf044362f5"
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "#@markdown Execute this cell to load data, create stimuli, and plot neural tuning curves\n",
    "\n",
    "### Load data and bin at 8 degrees\n",
    "# responses are split into test and train\n",
    "resp_train, resp_test, stimuli = load_data_split()\n",
    "n_stimuli, n_neurons = resp_train.shape\n",
    "print('resp_train contains averaged responses of %i neurons to %i binned stimuli' % (n_neurons, n_stimuli))\n",
    "#print(resp_train.shape)\n",
    "\n",
    "# also make stimuli into images\n",
    "orientations = np.linspace(0, 360, 61)[:-1] - 180\n",
    "grating_stimuli = np.zeros((60, 12, 16), np.float32)\n",
    "for i, ori in enumerate(orientations):\n",
    "  grating_stimuli[i] = grating(ori, res=0.025)#[18:30, 24:40]\n",
    "\n",
    "grating_stimuli = torch.from_numpy(grating_stimuli)\n",
    "print('grating_stimuli contains 60 stimuli of size 12 x 16')\n",
    "\n",
    "# Visualize tuning curves\n",
    "fig, axs = plt.subplots(3, 5, figsize=(15,7))\n",
    "for k, ax in enumerate(axs.flatten()):\n",
    "  neuron_index = np.random.choice(n_neurons)  # pick random neuron\n",
    "  plot_tuning(ax, stimuli, resp_train[:, neuron_index], resp_test[:, neuron_index], neuron_index, linewidth=2)\n",
    "  if k==0:\n",
    "    ax.text(1.0, 0.9, 'train', color='y', transform=ax.transAxes)\n",
    "    ax.text(1.0, 0.65, 'test', color='m', transform=ax.transAxes)\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bonus Section 3.2: Adding a fully-connected layer to create encoding model\n",
    "\n",
    "We will build a torch model like above with a convolutional layer. Additionally, we will add a fully connected linear layer from the convolutional units to the neurons. We will use 6 convolutional channels ($C^{out}$) and a kernel size ($K$) of 7 with a stride of 1 and padding of $K/2$ (same as above). The stimulus is size `(12, 16)`. Then the convolutional unit activations will go through a linear layer to be transformed into neural responses.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Think! 3.2: Number of unis and weights\n",
    "\n",
    "* How many units will the convolutional layer have?\n",
    "* How many weights will the fully connected linear layer from convolutional units to neurons have?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:51.743373Z",
     "iopub.status.busy": "2021-05-31T20:48:51.742752Z",
     "iopub.status.idle": "2021-05-31T20:48:51.746846Z",
     "shell.execute_reply": "2021-05-31T20:48:51.746283Z"
    }
   },
   "outputs": [],
   "source": [
    "# to_remove explanation\n",
    "\n",
    "\"\"\"\n",
    "1. There are 6 convolutional channels and 12 x 16 positions because the stride\n",
    "of the convolution is 1. Therefore, there are 6 * 12 * 16 = 1152 units in the\n",
    "convolutional layer.\n",
    "2. The fully connected linear layer will have 1152 * n_neurons weights. It will\n",
    "also have n_neurons additive bias terms.\n",
    "\"\"\";"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Bonus Coding Exercise 3.2: Add linear layer\n",
    "\n",
    "Remember in Tutorial 1 we used linear layers. Use your knowledge from Tutorial 1 to add a linear layer to the model we created above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:51.757820Z",
     "iopub.status.busy": "2021-05-31T20:48:51.754319Z",
     "iopub.status.idle": "2021-05-31T20:48:51.864631Z",
     "shell.execute_reply": "2021-05-31T20:48:51.865452Z"
    }
   },
   "outputs": [],
   "source": [
    "class ConvFC(nn.Module):\n",
    "  \"\"\"Deep network with one convolutional layer + one fully connected layer\n",
    "\n",
    "  Attributes:\n",
    "    conv (nn.Conv1d): convolutional layer\n",
    "    dims (tuple): shape of convolutional layer output\n",
    "    out_layer (nn.Linear): linear layer\n",
    "\n",
    "  \"\"\"\n",
    "\n",
    "  def __init__(self, n_neurons, c_in=1, c_out=6, K=7, b=12*16,\n",
    "               filters=None):\n",
    "    \"\"\" initialize layer\n",
    "    Args:\n",
    "        c_in: number of input stimulus channels\n",
    "        c_out: number of convolutional channels\n",
    "        K: size of each convolutional filter\n",
    "        h: number of stimulus bins, n_bins\n",
    "    \"\"\"\n",
    "    super().__init__()\n",
    "    self.conv = nn.Conv2d(c_in, c_out, kernel_size=K,\n",
    "                          padding=K//2, stride=1)\n",
    "    self.dims = (c_out, b)  # dimensions of conv layer output\n",
    "    M = np.prod(self.dims) # number of hidden units\n",
    "\n",
    "    ################################################################################\n",
    "    ## TO DO for students: add fully connected layer to network (self.out_layer)\n",
    "    # Fill out function and remove\n",
    "    raise NotImplementedError(\"Student exercise: add fully connected layer to initialize network\")\n",
    "    ################################################################################\n",
    "    self.out_layer = nn.Linear(M, ...)\n",
    "\n",
    "    # initialize weights\n",
    "    if filters is not None:\n",
    "      self.conv.weight = nn.Parameter(filters)\n",
    "      self.conv.bias = nn.Parameter(torch.zeros((c_out,), dtype=torch.float32))\n",
    "\n",
    "    nn.init.normal_(self.out_layer.weight, std=0.01) # initialize weights to be small\n",
    "\n",
    "  def forward(self, s):\n",
    "    \"\"\" Predict neural responses to stimuli s\n",
    "\n",
    "    Args:\n",
    "        s (torch.Tensor): p x L tensor with stimuli\n",
    "\n",
    "    Returns:\n",
    "        torch.Tensor: p x N tensor with convolutional layer unit activations.\n",
    "\n",
    "    \"\"\"\n",
    "    s = s.unsqueeze(1)  # p x 1 x W x H, add a singleton dimension for the single channel\n",
    "    a = self.conv(s)  # output of convolutional layer\n",
    "    a = a.view(-1, np.prod(self.dims))  # flatten each convolutional layer output into a vector\n",
    "\n",
    "    ################################################################################\n",
    "    ## TO DO for students: add fully connected layer to forward pass of network (self.out_layer)\n",
    "    # Fill out function and remove\n",
    "    raise NotImplementedError(\"Student exercise: add fully connected layer to network\")\n",
    "    ################################################################################\n",
    "    y = ...\n",
    "\n",
    "    return y\n",
    "\n",
    "\n",
    "device = torch.device('cpu')\n",
    "\n",
    "# (Optional) To speed up processing, go to \"Runtime\" menu and \"Change runtime\"\n",
    "# and select GPU processing, then uncomment line below, otherwise runtime will\n",
    "# be ~ 2 minutes\n",
    "# device = torch.device('cuda')\n",
    "\n",
    "# Initialize network\n",
    "n_neurons = resp_train.shape[1]\n",
    "K = 7\n",
    "# we will initialize with the same filters as above\n",
    "net = ConvFC(n_neurons, filters=filters(K))\n",
    "net = net.to(device)\n",
    "\n",
    "# Run GD on training set data\n",
    "# ** this time we are also providing the test data to estimate the test loss\n",
    "train_loss, test_loss = train(net, regularized_MSE_loss,\n",
    "                              train_data=grating_stimuli.to(device), train_labels=resp_train.to(device),\n",
    "                              test_data=grating_stimuli.to(device), test_labels=resp_test.to(device),\n",
    "                              n_iter=200, learning_rate=2,\n",
    "                              L2_penalty=5e-4, L1_penalty=1e-6)\n",
    "\n",
    "# Plot the training loss over iterations of GD\n",
    "plot_training_curves(train_loss, test_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-31T20:48:51.876717Z",
     "iopub.status.busy": "2021-05-31T20:48:51.876112Z",
     "iopub.status.idle": "2021-05-31T20:49:59.573023Z",
     "shell.execute_reply": "2021-05-31T20:49:59.572508Z"
    }
   },
   "outputs": [],
   "source": [
    "# to_remove solution\n",
    "class ConvFC(nn.Module):\n",
    "  \"\"\"Deep network with one convolutional layer + one fully connected layer\n",
    "\n",
    "  Attributes:\n",
    "    conv (nn.Conv1d): convolutional layer\n",
    "    dims (tuple): shape of convolutional layer output\n",
    "    out_layer (nn.Linear): linear layer\n",
    "\n",
    "  \"\"\"\n",
    "\n",
    "  def __init__(self, n_neurons, c_in=1, c_out=6, K=7, b=12*16,\n",
    "               filters=None):\n",
    "    \"\"\" initialize layer\n",
    "    Args:\n",
    "        c_in: number of input stimulus channels\n",
    "        c_out: number of convolutional channels\n",
    "        K: size of each convolutional filter\n",
    "        h: number of stimulus bins, n_bins\n",
    "    \"\"\"\n",
    "    super().__init__()\n",
    "    self.conv = nn.Conv2d(c_in, c_out, kernel_size=K,\n",
    "                          padding=K//2, stride=1)\n",
    "    self.dims = (c_out, b)  # dimensions of conv layer output\n",
    "    M = np.prod(self.dims) # number of hidden units\n",
    "    self.out_layer = nn.Linear(M, n_neurons)\n",
    "\n",
    "    # initialize weights\n",
    "    if filters is not None:\n",
    "      self.conv.weight = nn.Parameter(filters)\n",
    "      self.conv.bias = nn.Parameter(torch.zeros((c_out,), dtype=torch.float32))\n",
    "\n",
    "    nn.init.normal_(self.out_layer.weight, std=0.01) # initialize weights to be small\n",
    "\n",
    "  def forward(self, s):\n",
    "    \"\"\" Predict neural responses to stimuli s\n",
    "\n",
    "    Args:\n",
    "        s (torch.Tensor): p x L tensor with stimuli\n",
    "\n",
    "    Returns:\n",
    "        torch.Tensor: p x N tensor with convolutional layer unit activations.\n",
    "\n",
    "    \"\"\"\n",
    "    s = s.unsqueeze(1)  # p x 1 x W x H, add a singleton dimension for the single channel\n",
    "    a = self.conv(s)  # output of convolutional layer\n",
    "    a = a.view(-1, np.prod(self.dims))  # flatten each convolutional layer output into a vector\n",
    "    y = self.out_layer(a)\n",
    "    return y\n",
    "\n",
    "\n",
    "device = torch.device('cpu')\n",
    "\n",
    "# (Optional) To speed up processing, go to \"Runtime\" menu and \"Change runtime\"\n",
    "# and select GPU processing, then uncomment line below, otherwise runtime will\n",
    "# be ~ 2 minutes\n",
    "# device = torch.device('cuda')\n",
    "\n",
    "# Initialize network\n",
    "n_neurons = resp_train.shape[1]\n",
    "K = 7\n",
    "# we will initialize with the same filters as above\n",
    "net = ConvFC(n_neurons, filters=filters(K))\n",
    "net = net.to(device)\n",
    "\n",
    "# Run GD on training set data\n",
    "# ** this time we are also providing the test data to estimate the test loss\n",
    "train_loss, test_loss = train(net, regularized_MSE_loss,\n",
    "                              train_data=grating_stimuli.to(device), train_labels=resp_train.to(device),\n",
    "                              test_data=grating_stimuli.to(device), test_labels=resp_test.to(device),\n",
    "                              n_iter=200, learning_rate=2,\n",
    "                              L2_penalty=5e-4, L1_penalty=1e-6)\n",
    "\n",
    "# Plot the training loss over iterations of GD\n",
    "with plt.xkcd():\n",
    "  plot_training_curves(train_loss, test_loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How well can we fit single neuron tuning curves with this model? What aspects of the tuning curves are we capturing?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 430
    },
    "execution": {
     "iopub.execute_input": "2021-05-31T20:49:59.581649Z",
     "iopub.status.busy": "2021-05-31T20:49:59.579736Z",
     "iopub.status.idle": "2021-05-31T20:50:00.561417Z",
     "shell.execute_reply": "2021-05-31T20:50:00.561818Z"
    },
    "outputId": "b2bb7c5f-39da-4e61-bea0-aeeb9338c015"
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "#@markdown Execute this cell to examine predictions for random subsets of neurons\n",
    "\n",
    "y_pred = net(grating_stimuli.to(device))\n",
    "# Visualize tuning curves & plot neural predictions\n",
    "fig, axs = plt.subplots(2, 5, figsize=(15,6))\n",
    "for k, ax in enumerate(axs.flatten()):\n",
    "  ineur = np.random.choice(n_neurons)\n",
    "  plot_prediction(ax, y_pred[:,ineur].detach().cpu(),\n",
    "                  resp_train[:,ineur],\n",
    "                  resp_test[:,ineur])\n",
    "  if k==0:\n",
    "    ax.text(.1, 1., 'train', color='y', transform=ax.transAxes)\n",
    "    ax.text(.1, 0.9, 'test', color='m', transform=ax.transAxes)\n",
    "    ax.text(.1, 0.8, 'prediction', color='g', transform=ax.transAxes)\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see if the convolutional channels changed at all from their initialization as center-surround and Gabor filters. If they don't then it means that they were a sufficient basis set to explain the responses of the neurons to orientations to the accuracy seen above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 283
    },
    "execution": {
     "iopub.execute_input": "2021-05-31T20:50:00.566401Z",
     "iopub.status.busy": "2021-05-31T20:50:00.565938Z",
     "iopub.status.idle": "2021-05-31T20:50:01.094405Z",
     "shell.execute_reply": "2021-05-31T20:50:01.093882Z"
    },
    "outputId": "607f4c80-4b7b-4c6f-b442-293c0ac1ae16"
   },
   "outputs": [],
   "source": [
    "# get weights of conv layer in convLayer\n",
    "weights = net.conv.weight.detach().cpu()\n",
    "print(weights.shape) # can you identify what each of the dimensions are?\n",
    "\n",
    "plot_weights(weights, channels=np.arange(0, conv_channels))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "W2D1_Tutorial2",
   "provenance": [],
   "toc_visible": true
  },
  "kernel": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
